#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2019/11/2 11:36
# @Author  : CHEN Wang
# @Site    :
# @File    : db_conn.py
# @Software: PyCharm

"""
脚本说明: 创建mysql数据库连接， 对mysql数据库进行相关的操作
"""

import math
import datetime
import time
import pymysql
import pandas as pd
from sqlalchemy import create_engine
from quant_researcher.quant.project_tool.logger.my_logger import LOG
from quant_researcher.quant.project_tool.conf import LOCAL_CONF
from quant_researcher.quant.project_tool import time_tool


def if_this_is_a_p_key_of_the_table(target_p_key, table, conn):
    unique_key_dict = get_the_primary_key_by_table(table, conn=conn)
    if unique_key_dict:
        for a_p_key in unique_key_dict.values():
            # print('set(a_p_key) - set(target_p_key):', set(a_p_key) - set(target_p_key))
            # print('set(target_p_key) - set(a_p_key):', set(target_p_key) - set(a_p_key))
            if set(a_p_key) - set(target_p_key) == set() == set(target_p_key) - set(a_p_key):
                return True
    LOG.warning(f'表 {table} 的主键中不含有 {target_p_key}，只找到：{unique_key_dict}')
    return False


def get_the_primary_key_by_table(x, conn):
    sql = f'show index from `{x}`'
    col_key_name = 'Key_name'
    col_non_unique = 'Non_unique'
    col_index_type = 'Index_type'
    col_nullable = 'Null'
    col_cardinality = 'Cardinality'
    col_column_name = 'Column_name'
    col_seq_in_index = 'Seq_in_index'
    df = execute(sql=sql, need_4_df=True, conn=conn)
    df = df[df[col_non_unique] == 0]  # 找到这个表定义了的唯一主键，可能多个
    key_names = list(df.drop_duplicates(subset=col_key_name, keep='first')[col_key_name].values)
    unique_key_dict = {}
    LOG.info(f'索引：{key_names}')
    for a_key in key_names:
        tmp_df = df[df[col_key_name] == a_key].copy()
        tmp_df.sort_values(by=col_seq_in_index, ascending=True, inplace=True)
        nullable = tmp_df[col_nullable].values[0]
        cardinality = tmp_df[col_cardinality].values[0]
        index_type = tmp_df[col_index_type].values[0]
        cols_included = list(tmp_df[col_column_name].values)
        LOG.info(f'索引={a_key}，可空={nullable}，类型={index_type}，基数={cardinality}，包含的列(按顺序)={cols_included}')
        unique_key_dict.update({a_key: cols_included})
    LOG.info(f'表 {x} 共有{len(unique_key_dict)}个唯一主键！')

    return unique_key_dict


def on_duplicate_key_update(a_lst, primary_key, conn, table, **kwargs):
    """
    批量更新
        无则插入，有则更新对应字段，类似如下SQL：
            insert into table () values () on duplicate key update
    这种方法将去检查这个表的唯一索引，如果插入会引起主键冲突则更新对应字段，否则插入
    那么就必须保证这个的索引的维护，如果表的索引不是预期的，那么将会出现重复的记录
        警告：通常，您应该尽量避免对带有多个唯一关键字的表使用ON DUPLICATE KEY子句

    :param list_OR_df a_lst: a list of dict 或 pd.DataFrame，传入 df 时不会影响原数据
    :param list primary_key: str or list of str，类似：['fund_code', 'p_count']
    :param conn conn: 数据库连接
    :param str table: 表
    :param kwargs:
        - create_missing_cols，bool，传入的 list of dict，可能键不对齐，设置此参数来自动对齐
            默认： False
        - replace_nan，bool，np.nan有时不能直接写到 MySQL，设置此参数来自动转成 None
            默认： False
        - create_time_cols,，bool，是否要增加titime、tutime、tutimestamp这三列
            默认： True
    :return:
    """
    updated_on_ts_col = 'tutimestamp'
    inserted_on_col, updated_on_col = 'titime', 'tutime'

    create_missing_cols = kwargs.pop('create_missing_cols', False)
    if create_missing_cols:
        # 传来的一个list，里面如果每个dict中的包含的键有差别
        # 可通过先创建DF，再转，使所有的dict包含相同的键
        if isinstance(a_lst, list):
            a_lst = pd.DataFrame(a_lst)
            a_lst = [dict(x[1]) for x in a_lst.iterrows()]
        else:
            LOG.warning(f'要像自动创建缺失列，却传的不是list（{a_lst.__class__}）。'
                        '请检查输入对象类型！')

    if isinstance(a_lst, pd.DataFrame):
        # TODO 暂时不太确定，还没决定这种 多入参类型 的方式的好坏
        a_lst = [dict(x[1]) for x in a_lst.iterrows()]  # 感觉这是99.9999%场景
    if not a_lst:
        LOG.warning(f'传入了假数组：{a_lst}')
        return
    if isinstance(primary_key, tuple):
        primary_key = [primary_key]
    if not if_this_is_a_p_key_of_the_table(target_p_key=primary_key,
                                           table=table, conn=conn):
        raise RuntimeError('目标表主键与所提供不匹配')

    replace_nan = kwargs.pop('replace_nan', False)
    if replace_nan:
        for y in a_lst:
            for k, value in y.items():
                if value is not None \
                        and not isinstance(value, str) \
                        and not isinstance(value, datetime.date) \
                        and (math.isnan(value) or math.isinf(value)):
                    y[k] = None

    create_time_cols = kwargs.pop('create_time_cols', True)
    i_t, u_t, u_t_ts = time_tool.get_datetime_and_timestamp()
    if create_time_cols:
        for x in a_lst:
            x.update({
                inserted_on_col: i_t, updated_on_col: u_t, updated_on_ts_col: u_t_ts
            })

    v_data = [tuple(x.values()) for x in a_lst]
    # TODO 如果每一个字典的键的顺序不一样怎么办？
    #  你就这样直接转 tuple，直接取第一个字典的键的顺序，是不对的！！！
    #  zhouqiang，2020年1月16日17:55:30
    m_keys = a_lst[0].keys()

    cols_found = ', '.join([f'`{xx}`' for xx in m_keys])
    cols_2_set = ', '.join(['%s' for _ in m_keys])
    cols_2_update = []
    for a_col in m_keys:
        if a_col not in primary_key and a_col != inserted_on_col:
            cols_2_update.append(f'`{a_col}`=values(`{a_col}`)')
    cols_2_update = ', '.join(cols_2_update)
    sql = f'insert into {table} ' \
          f'({cols_found}) values ({cols_2_set}) ' \
          f'on duplicate key update {cols_2_update}'
    LOG.info(f'最近更新 {table}（主键：{primary_key}）发生在：{u_t}')
    batch_update(sql=sql, values_list=v_data, conn=conn)


def batch_update(sql, values_list, **kwargs):
    """
    对 MySQL 的表进行批量更新

    :param str sql: 批量执行的 sql
    :param list values_list: sql 中的值将由 list 中的 tuple 按顺序填充
    :param kwargs:
        - conn，数据库连接，默认：None（从配置读取并创建）
    :return: None
    """
    value_num = len(values_list)
    if value_num == 0:
        LOG.warning('传入了空数组，请检查！')
        return
    LOG.info(f'批量更新（{value_num}行）：\n      %s', sql)
    conn = kwargs.pop('conn', None)
    close_it = False
    if conn is None:
        conn = get_mysql_conn()
        close_it = True
    try:
        cursor = conn.cursor()
        effect_row = cursor.executemany(sql, values_list)
        LOG.info('影响行数⭐：%s', effect_row)
        conn.commit()  # 提交,数据才被真正写到了数据库中
    except Exception as err:
        conn.rollback()  # 回滚操作
        LOG.error(f'执行：\n{sql}\n的过程中发生错误：{str(err)}')
        for a_value in values_list[:199]:
            LOG.error(f'引起错误的数据：{a_value}')
        raise err
    finally:
        if close_it:
            conn.close()


def execute(sql, **kwargs):
    """
    简单地执行一个写好的sql

    :param str sql: 需要执行的sql
    :param kwargs:
        - need_4_df，bool，这句话是不是想要获得一个DF，默认：False
        - conn，connection对象，默认：自动获取，根据配置
    :return:
    """
    need_4_df = kwargs.pop('need_4_df', False)
    conn = kwargs.get('conn')

    close_it = False
    if conn is None:
        conn = get_mysql_conn()
        close_it = True
    try:
        LOG.info(f'开始执行SQL(是否返回DF：{need_4_df}) [{sql}] ...')
        if need_4_df:
            df = pd.read_sql_query(sql, conn)
            if len(sql) < 170:  # 根据我的显示器尺寸得到的
                LOG.info('行：%s → SQL：%s', len(df), sql)
            else:
                LOG.info('行：%s → SQL：↓\n     %s', len(df), sql)
            return df
        else:
            cursor = conn.cursor()
            effect_row = cursor.execute(sql)
            LOG.info(f'⭐影响行数⭐：⭐{effect_row}⭐ → SQL：↓\n     {sql}')
            conn.commit()  # 提交,数据才被真正写到了数据库中
    except Exception as err:
        LOG.error(f'☠ 执行SQL发生错误：{err} ☠ → SQL：↓\n     {sql}')
        conn.rollback()  # 回滚操作
        raise err
    finally:
        if close_it:
            conn.close()


def delete_from(table, conn, where=None):
    if where is None:
        sql = f'truncate table {table}'
    else:
        if not isinstance(where, list):
            raise ValueError(f'where应该是list，却传入：{where.__class__}')
        condition = [f'({x})' for x in where]
        condition = ' and '.join(condition)
        sql = f'delete from {table} where {condition}'
    execute(sql=sql, conn=conn)


def read_v2(db_str=None, table_str=None, select=None, sfrom=None, where=None, conn=None, order_by=None, limit=None,
            **kwargs):
    """
    查询数据，最优实践，2019年4月26日14:29:42

    2019年4月29日09:40:24
    一个项目，或者一个开发环境，同时使用多个MySQL数据库的例子，还是少
    一个Python脚本，运行时间一般都不长，如果很长的话，就应该考虑其他的问题了
    不管运行一个package中的哪一个脚本，只要是import过的东西，都会在进程中存在，直到进程结束
    因此，应该在这里定义一个conn，从第一次（实际用到的时候生成）生成，后面应该重复使用，进程结束之后，关闭（或者由数据库自己删除无效连接）
    但是，如果这里不手动close，建立的这个连接在数据库端，完全靠设置的连接超时等待时间决定，如果那个时间设置的有点大，在一段时间内，这个脚本又被
    调用了很多次，在调试或者数据重算，bug等情况下，可能会引起数据库那边连接出现问题，因为那边有太多的等待中的连接，可能导致无法创建新的连接，
    对其他用户产生影响，但是，这些应该都算是特殊情况或者bug出现，so，暂时先不考虑吧
    说明：
        → 当然，这个函数默认使用已经生成的conn，如果想使用其他的连接，在特殊情况下，也支持传入conn对象
        → 为了不影响之前的代码或者以后出现变化，这里采用一个装饰器，如果关键字参数中如果没有，就通过装饰器加入，若有了，便不管了
        → 至于装饰器，好像最好是使用参数名传递这种方式写代码，通过位置的话，就包含了不管是代码编写者还是阅读者都对那个函数很了解了，
            但是事实好像并没有

    :param db_str: 数据库名，与conn两者，有一便可，首先判断是不是有 conn，如果没有，将会使用这个db建立连接
    :param table_str: 表名，与sfrom两者，有一便可，首先判断是不是有 table，如果没有，将会使用这个table建立连接
    :param select: string或者list of string，需要哪些列，默认：*
    :param sfrom: sql的from后跟的东西，为了适配某些特殊场景，与前面的table做了区分，如果没传的话，使用table，默认：None
    :param where: string或者list of string，查询的条件，默认：None
    :param conn: 查询使用的连接对象，为了适应某些频繁查询或其他场景，使用外部的连接对象，默认：None
    :param order_by: 如果要排序，默认：None
    :param limit: 返回的最多条数，默认：None
    :param kwargs:
        - where_and，bool，若 where 是 list，用 and 还是 or 来连接，默认：True(and)
    :return: df，DataFrame
    """
    # engine_str = '%s/%s' % (LOCAL_CONF['mysql']['engine'], db_str)
    # LOG.info('引擎字符串：%s。', engine_str)
    # engine = create_engine(engine_str)
    where_and = kwargs.pop('where_and', True)
    tmp_str = table_str if sfrom is None else sfrom
    LOG.info(f'表：{tmp_str}，排序：{order_by}，LIMIT：{limit}'
             f'\n  SELECT：{select}\n   WHERE：{where}')
    close_conn = False
    if conn is None:
        if db_str is None:
            db_str = LOCAL_CONF['mysql']['default_db']
        LOG.info(f'未传入conn，将建立连接，DB={db_str}')
        info = {
            'database': db_str,
            'host': LOCAL_CONF['mysql']['ip'],
            'port': LOCAL_CONF.getint('mysql', 'port'),
            'user': LOCAL_CONF['mysql']['user'],
            'password': LOCAL_CONF['mysql']['password']
        }
        conn = pymysql.connect(**info)
        close_conn = True  # 如果是在这里自己创建的conn，那么用完了，需要close，若传进来的就不管了

    if select is None:
        cols = '*'
    elif isinstance(select, str):
        cols = select
    elif isinstance(select, list):
        cols = ', '.join(select)
    else:
        raise RuntimeError

    if order_by is None:
        pass
    elif isinstance(order_by, str):
        pass
    elif isinstance(order_by, list):
        order_by = ', '.join(order_by)
    else:
        raise RuntimeError

    if sfrom is None:
        sfrom = table_str
    if sfrom is None:
        raise ValueError(f'sfrom 和 表名都没传，不能这样！')

    if where is None:
        condition = ''
    elif isinstance(where, str):
        condition = ' where %s' % where
    elif isinstance(where, list):
        if where_and:
            tmp = ' and '
        else:
            tmp = ' or '
        condition = tmp.join(map(lambda x: '(%s)' % x, where))
        if condition:
            condition = ' where %s' % condition
    else:
        raise RuntimeError
    sql_str = 'select %s from %s%s' % (cols, sfrom, condition)

    if order_by is not None:
        sql_str += ' order by %s' % order_by

    if limit is not None:
        sql_str += ' limit %s' % limit

    # 2019年5月13日13:54:39 这里的日志总感觉很别扭，而且也没有catch exception，还是再改改
    try:
        # df = pd.read_sql_query(sql_str, engine)
        df = pd.read_sql_query(sql_str, conn)
        LOG.info(f'囍：{len(df)}，SQL：' + ('↓\n  ' if len(sql_str) > 170 else ' → ') + f'{sql_str}')  # 170，适应显示器
        return df
    except Exception as err:
        LOG.error(f'☠ 执行SQL发生错误({err})↓\n  SQL：{sql_str}')
        raise err
    finally:
        if close_conn:
            LOG.info('Going to ♻ this connection!!!')  # 🔫
            conn.close()


def read_v3(sql_str, conn):
    """
    一般的sql语句，都可以通过read_v2从数据库中读数据
    但是，如果sql语句比较复杂，比如涉及关联子查询，read_v2无法处理这种复杂的sql语句
    对于这种复杂的sql语句，直接用pd.read_sql读取，然后将sql语句通过日志输出

    :param str sql_str: 需要执行的sql语句
    :param conn: 数据库连接
    :return: pd.DataFrame
    """
    try:
        df = pd.read_sql_query(sql_str, conn)
        LOG.info(f'囍：{len(df)}，SQL：' + ('↓\n  ' if len(sql_str) > 170 else ' → ') + f'{sql_str}')  # 170，适应显示器
        return df
    except Exception as err:
        LOG.error(f'☠ 执行SQL发生错误({err})↓\n  SQL：{sql_str}')
        raise err


#####################################   创建mysql连接  #############################################

def get_mysql_conn_by_name(x, sector=None, cn_name=None, **kwargs):
    """

    :param x: 数据库名字前缀
    :param sector: 数据库分类
    :param cn_name: 数据库中文名
    :param kwargs:
        - bool by_sqlalchemy: 是否是通过sqlalchemy.creat_engine创建，还是通过pymysql的方式创建，
                              默认用pymysql， 可以通过on_duplicate_key_update函数写入
                              通过creat_engine创建写入时，不方便（不能通过on_duplicate_key_update
                              函数写入，只能通过to_sql写入）。但是创建的连接不容易断掉
    :return:
    """

    if sector is None:
        sector = LOCAL_CONF['other']['default_section']
    database = LOCAL_CONF[sector][f'{x}_db']
    host = LOCAL_CONF[sector][f'{x}_host']
    user = LOCAL_CONF[sector][f'{x}_user']
    password = LOCAL_CONF[sector][f'{x}_password']
    port = LOCAL_CONF.getint(sector, f'{x}_port')
    msg = f'DB={x}-Of-{sector.upper()}' if cn_name is None else cn_name
    by_sqlalchemy = kwargs.pop('by_sqlalchemy', False)

    conn = get_mysql_conn(database=database, host=host, user=user, password=password, port=port,
                          msg=msg, by_sqlalchemy=by_sqlalchemy)
    return conn


def get_mysql_conn(**kwargs):
    """
    构造一个数据库连接对象

    :param kwargs:
        - host，str，主机IP，默认：从本地配置文件中读取
        - port，int，端口，默认：从本地配置文件中读取
        - user，str，用户名，默认：从本地配置文件中读取
        - password，str，密码，默认：从本地配置文件中读取
        - database，str，数据库，默认：从本地配置文件中读取
        - msg，str，附加信息，默认：未提供
        - retry_times，int，尝试连接重试次数，默认：1
        - retry_after，float，尝试连接重试等待时间（秒），默认：1.5
    :return: 连接对象
    """
    host = kwargs.pop('host', LOCAL_CONF['mysql']['ip'])
    port = kwargs.pop('port', LOCAL_CONF['mysql']['port'])
    user = kwargs.pop('user', LOCAL_CONF['mysql']['user'])
    password = kwargs.pop('password', LOCAL_CONF['mysql']['password'])
    database = kwargs.pop('database', LOCAL_CONF['mysql']['default_db'])
    msg = kwargs.pop('msg', '未提供')
    by_sqlalchemy = kwargs.pop('by_sqlalchemy', False)
    retry_times = kwargs.pop('retry_times', 1)
    retry_after = kwargs.pop('retry_after', 1.5)

    msg = f'（DB：{msg}<{user}@{host}:{port}/{database}>）'
    LOG.warning(f'嘗試連接{msg}，重試（次數={retry_times}，間隔={retry_after}秒）')
    ii = 1
    while ii <= retry_times:
        try:
            if by_sqlalchemy:
                LOG.info(f'通过sqlalchemy创建数据库连接')
                conn = create_engine(f"mysql+pymysql://{user}:{password}@{host}:{port}/{database}")
            else:
                LOG.info(f'通过pymysql创建数据库连接')
                conn = pymysql.connect(host=host, port=port, user=user, password=password,
                                       database=database)
            LOG.warning(f'用完連接{msg}，最好手動關閉')
            return conn
        except Exception as err:
            LOG.warning(f'第 {ii} 次嘗試連接{msg}失敗，原因：{err}')
            if ii == retry_times:
                LOG.error(f'{retry_times} 次嘗試之後，還是連不上{msg}')
                raise err
            time.sleep(retry_after)
        ii += 1
